package com.hongji.lms5xxcommunicator.realtime.websoket;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
@Data
public class StreamHub {

    private final Set<WebSocketSession> sessions = ConcurrentHashMap.newKeySet();

    /** 老字段：首片 full 帧（兜底用） */
    private volatile String lastFullFrameJson;

    /** 新字段：最近一次整帧的“所有分片” */
    private volatile List<String> lastFullFrameChunks; // 不可变列表，由 broadcaster 设置

    /** 网格规格（gridSpec） */
    private volatile String lastGridSpecJson;

    /** 新会话加入：先规格，再整帧（分片逐条），最后兜底发首片 */
    void add(WebSocketSession s) {
        sessions.add(s);
        log.info("[WS] add session: {}, total={}", s.getId(), sessions.size());

        try {
            if (lastGridSpecJson != null && s.isOpen()) {
                s.sendMessage(new TextMessage(lastGridSpecJson));
                log.debug("[WS] replay gridSpec to {}", s.getId());
            }
        } catch (IOException e) {
            log.warn("[WS] replay gridSpec fail: {}, {}", s.getId(), e.getMessage());
        }

        // ★ 优先重放“整帧的所有分片”
        List<String> chunks = lastFullFrameChunks;
        if (chunks != null && !chunks.isEmpty()) {
            for (String part : chunks) {
                try {
                    if (s.isOpen()) s.sendMessage(new TextMessage(part));
                } catch (IOException e) {
                    log.warn("[WS] replay chunk fail: {}, {}", s.getId(), e.getMessage());
                }
            }
            log.debug("[WS] replay full frame chunks: {} -> {}", chunks.size(), s.getId());
            return; // 已完整回放，无需再发首片
        }

        // 兜底：只有首片缓存时
        try {
            if (lastFullFrameJson != null && s.isOpen()) {
                s.sendMessage(new TextMessage(lastFullFrameJson));
                log.debug("[WS] replay full frame (single) to {}", s.getId());
            }
        } catch (IOException e) {
            log.warn("[WS] replay full frame fail: {}, {}", s.getId(), e.getMessage());
        }
    }

    void remove(WebSocketSession s) {
        sessions.remove(s);
        log.info("[WS] remove session: {}, total={}", s.getId(), sessions.size());
    }

    public void broadcast(String text) {
        TextMessage msg = new TextMessage(text);
        int ok = 0, fail = 0;
        for (WebSocketSession s : sessions) {
            if (s.isOpen()) {
                try { s.sendMessage(msg); ok++; }
                catch (IOException e) { fail++; log.warn("[WS] send fail: {}, error={}", s.getId(), e.getMessage()); }
            }
        }
        log.info("[WS] broadcast done: ok={}, fail={}, total={}", ok, fail, sessions.size());
    }

    public int size() { return sessions.size(); }
}
